import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号

# 定义函数f(z)
def f(z):
    # doubao
    # return (np.sqrt(1j*(z - 1)) - 1) / (np.sqrt(1j*(z - 1)) + 1)

    # tongyi
    # return (np.sqrt(1 - z) - 1j) / (np.sqrt(1 - z) + 1j)

    # deepseek
    # return 2/np.sqrt(2-z)-1
    
    # return 1 - 2/z**2

    return 1 - 2/np.sqrt(z)

# 创建z平面网格点
x = np.linspace(-6, 6, 400)
y = np.linspace(-6, 6, 400)
X, Y = np.meshgrid(x, y)
Z = X + 1j * Y

# 计算抛物线y^2=4(1-x)左边部分的掩码
mask = Y**2 > 4 * (1 - X)

# 计算双曲线y^2=x^2-1 右边部分的掩码
# mask = Y**2 < X**2 - 1

# 应用变换到w平面上
W = np.zeros_like(Z)
W[mask] = f(Z[mask])

# 绘制图像
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))

# 在z平面上绘制原始区域
ax1.plot(X[mask], Y[mask], 'b.', markersize=1)
ax1.set_title('z-plane: $y^2 > 4(1-x)$')
ax1.set_xlabel('Re(z)')
ax1.set_ylabel('Im(z)')
ax1.set_xlim([-6,6])
ax1.set_ylim([-6,6])
ax1.grid(True)

# 在w平面上绘制映射后的区域
ax2.plot(np.real(W[mask]), np.imag(W[mask]), 'g.', markersize=1)
ax2.set_title('w-plane: $w=1-2/\sqrt{z}$, $u^2+v^2<1$')
ax2.set_xlabel('Re(w)')
ax2.set_ylabel('Im(w)')
ax2.set_xlim([-2,2])
ax2.set_ylim([-2,2])
ax2.grid(True)

plt.tight_layout()
# plt.show()